/*-------------------------------------------------------------------------
 * Copyright (c) Microsoft Corporation.  All rights reserved.
 *
 * src/service/tls.rs
 *
 *-------------------------------------------------------------------------
 */

//! TLS certificate management and SSL/TLS configuration for the `DocumentDB` gateway.
//!
//! This module provides comprehensive certificate management capabilities including
//! automatic certificate generation, file-based certificate loading, certificate
//! bundle management, and automatic certificate reloading for zero-downtime updates.
//!
//! # Features
//!
//! - **Certificate Store Management**: Handles both file-based and auto-generated certificates
//! - **Certificate Bundles**: Manages server certificates, private keys, and CA chains
//! - **TLS Provider**: Automatic certificate reloading with background monitoring
//! - **SSL Acceptor Configuration**: Integration with OpenSSL for secure connections
//! - **Zero-Downtime Updates**: Hot reloading of certificates without service interruption
//!
//! # Certificate Types
//!
//! The module supports two main certificate input types:
//! - `PemFile`: Load certificates from existing PEM files on disk
//! - `PemAutoGenerated`: Automatically generate self-signed certificates for development
//!
//! # Architecture
//!
//! ```text
//! ┌─────────────────┐    ┌──────────────────┐    ┌─────────────────┐
//! │ CertificateStore│───▶│ CertificateBundle│───▶│   TlsProvider   │
//! └─────────────────┘    └──────────────────┘    └─────────────────┘
//!          │                       │                       │
//!          ▼                       ▼                       ▼
//!   File Management        Certificate Loading      Automatic Reloading
//! ```

use std::{path::Path, sync::Arc};

use arc_swap::ArcSwap;
use openssl::{
    hash::{Hasher, MessageDigest},
    pkey::{Id, PKey, PKeyRef, Private},
    ssl::{SslAcceptor, SslAcceptorBuilder, SslCipherRef},
    x509::X509,
};

use crate::{
    configuration::{CertInputType, CertificateOptions},
    error::{DocumentDBError, Result},
    service::docdb_openssl,
};

/// Default key paths for auto-generated certificates
const DEFAULT_PRIVATE_KEY_PATH: &str = "./pkey.pem";
const DEFAULT_PUBLIC_KEY_PATH: &str = "./cert.pem";

/// Certificate file change monitoring interval in seconds
const CERT_FILES_CHANGE_WATCH_INTERVAL: u64 = 60;

/// Internal certificate store that manages certificate file paths.
///
/// The main difference between `CertificateStore` and `CertificateOptions` is that
/// `CertificateStore` will generate the keys (and paths to them) if `cert_type` is
/// `CertInputType::PemAutoGenerated`.
pub struct CertificateStore {
    certificate_path: String,
    private_key_path: String,
    ca_path: Option<String>,
}

impl CertificateStore {
    /// Creates a new certificate store from the provided options.
    ///
    /// For `PemFile` type, validates that required file paths are provided.
    /// For `PemAutoGenerated` type, generates certificates if they don't exist.
    ///
    /// # Arguments
    ///
    /// * `certificate_options` - Configuration specifying certificate type and paths
    ///
    /// # Returns
    ///
    /// Returns a configured `CertificateStore` ready for use.
    ///
    /// # Errors
    ///
    /// This function will return an error if:
    /// * Required file paths are missing for PemFile type
    /// * Certificate generation fails for auto-generation type
    /// * File system operations fail
    pub fn new(certificate_options: &CertificateOptions) -> Result<Self> {
        match certificate_options.cert_type {
            CertInputType::PemFile => {
                let certificate_path = certificate_options.file_path.as_ref().ok_or_else(|| {
                    DocumentDBError::internal_error(
                        "Certificate file path is required.".to_string(),
                    )
                })?;

                let private_key_path =
                    certificate_options.key_file_path.as_ref().ok_or_else(|| {
                        DocumentDBError::internal_error(
                            "Private key file path is required.".to_string(),
                        )
                    })?;

                Ok(Self {
                    certificate_path: certificate_path.clone(),
                    private_key_path: private_key_path.clone(),
                    ca_path: certificate_options.ca_path.clone(),
                })
            }
            CertInputType::PemAutoGenerated => {
                // We don't need to regenerate the keys if they already exist
                // One of the cases is e2e tests
                if !Path::new(DEFAULT_PRIVATE_KEY_PATH).exists()
                    || !Path::new(DEFAULT_PUBLIC_KEY_PATH).exists()
                {
                    docdb_openssl::generate_auth_keys(
                        DEFAULT_PRIVATE_KEY_PATH,
                        DEFAULT_PUBLIC_KEY_PATH,
                    )?;
                }

                Ok(Self {
                    certificate_path: DEFAULT_PUBLIC_KEY_PATH.to_string(),
                    private_key_path: DEFAULT_PRIVATE_KEY_PATH.to_string(),
                    ca_path: None,
                })
            }
        }
    }
}

/// A bundle containing SSL/TLS certificate, private key, and CA chain.
///
/// This struct holds all the cryptographic materials needed for SSL/TLS operations,
/// including the server certificate, its private key, and any intermediate CA certificates.
pub struct CertificateBundle {
    certificate: X509,
    private_key: PKey<Private>,
    ca_chain: Vec<X509>,
}

impl CertificateBundle {
    /// Creates a certificate bundle by loading certificates from the certificate store.
    ///
    /// This async function reads certificate files from disk and parses them into
    /// the appropriate OpenSSL structures.
    ///
    /// # Arguments
    ///
    /// * `certificate_store` - The certificate store containing file paths
    ///
    /// # Returns
    ///
    /// Returns a `CertificateBundle` on success.
    ///
    /// # Errors
    ///
    /// This function will return an error if:
    /// * Certificate or private key files cannot be read
    /// * Certificate parsing fails (invalid PEM format)
    /// * CA chain file is specified but cannot be read or parsed
    pub async fn from_cert_store(certificate_store: &CertificateStore) -> Result<Self> {
        let cert_bytes = tokio::fs::read(&certificate_store.certificate_path).await?;
        let certificate = X509::from_pem(&cert_bytes)?;

        let pkey_bytes = tokio::fs::read(&certificate_store.private_key_path).await?;
        let pkey = PKey::private_key_from_pem(&pkey_bytes)?;

        let mut ca_chain = Vec::new();
        if let Some(ca_path) = &certificate_store.ca_path {
            let ca_chain_pem = tokio::fs::read(ca_path).await?;
            ca_chain = X509::stack_from_pem(&ca_chain_pem)?;
        }

        Ok(Self {
            certificate,
            private_key: pkey,
            ca_chain,
        })
    }

    /// Returns a clone of the server certificate.
    ///
    /// Note: The OpenSSL crate doesn't provide borrowing for X509 certificates,
    /// so this method returns a clone.
    pub fn certificate(&self) -> X509 {
        self.certificate.clone()
    }

    /// Returns a reference to the private key.
    pub fn private_key(&self) -> &PKeyRef<Private> {
        self.private_key.as_ref()
    }

    /// Returns a slice of the CA certificate chain.
    pub fn ca_chain(&self) -> &[X509] {
        &self.ca_chain
    }
}

/// Creates an SSL acceptor builder configured with the specified certificate options.
///
/// This function is a convenience wrapper that combines certificate store creation,
/// certificate bundle loading, and TLS acceptor configuration into a single operation.
/// It handles both file-based certificates and auto-generated certificates based on
/// the provided configuration.
///
/// # Arguments
///
/// * `certificate_options` - Configuration specifying the certificate type, file paths,
///   and other SSL/TLS options
/// * `acceptor_builder` - Optional custom function for creating SSL acceptor builders
///
/// # Returns
///
/// Returns an `SslAcceptorBuilder` that can be further configured or built into
/// an `SslAcceptor` for accepting TLS connections.
///
/// # Errors
///
/// This function will return an error if:
/// * Certificate store creation fails (missing file paths, generation errors)
/// * Certificate files cannot be read or are in invalid PEM format
/// * Private key files cannot be read or parsed
/// * CA chain files are specified but cannot be loaded
/// * TLS acceptor configuration fails
pub async fn create_tls_acceptor_builder(
    certificate_options: &CertificateOptions,
    acceptor_builder: Option<docdb_openssl::AcceptorBuilderFn>,
) -> Result<SslAcceptorBuilder> {
    let cert_store = CertificateStore::new(certificate_options)?;
    let cert_bundle = CertificateBundle::from_cert_store(&cert_store).await?;

    docdb_openssl::create_tls_acceptor(&cert_bundle, acceptor_builder)
}

/// A provider that manages SSL/TLS certificates with automatic reloading.
///
/// This struct handles certificate loading, monitoring for file changes, and
/// automatic reloading when certificate files are updated on disk. It runs
/// a background task that periodically checks for certificate file modifications
/// every 60 seconds.
///
/// The provider uses `ArcSwap` to allow atomic updates of the SSL acceptor
/// without blocking ongoing connections.
#[derive(Clone)]
pub struct TlsProvider {
    tls_acceptor: Arc<ArcSwap<SslAcceptor>>,
    ciphersuite_mapping: Option<fn(Option<&str>) -> i32>,
    certificate_bundle: Arc<ArcSwap<CertificateBundle>>,
}

impl TlsProvider {
    /// Creates a new TLS provider with automatic certificate reloading.
    ///
    /// This function sets up the initial certificate bundle and starts a background
    /// task that monitors certificate files for changes and reloads them automatically.
    /// The background task runs every 60 seconds to check for certificate updates.
    ///
    /// # Arguments
    ///
    /// * `certificate_options` - Configuration specifying certificate type and paths
    /// * `acceptor_builder` - Optional custom function for creating SSL acceptor builders
    ///
    /// # Returns
    ///
    /// Returns a `TlsProvider` that manages SSL/TLS acceptors and automatic reloading.
    ///
    /// # Errors
    ///
    /// This function will return an error if:
    /// * Initial certificate loading fails
    /// * Certificate store creation fails
    /// * File modification time cannot be determined
    /// * TLS acceptor creation fails
    pub async fn new(
        certificate_options: &CertificateOptions,
        acceptor_builder: Option<docdb_openssl::AcceptorBuilderFn>,
        ciphersuite_mapping: Option<fn(Option<&str>) -> i32>,
    ) -> Result<Self> {
        let cert_store = CertificateStore::new(certificate_options)?;

        let tls_builder =
            create_tls_acceptor_builder(certificate_options, acceptor_builder).await?;
        let tls_acceptor_arc = Arc::new(ArcSwap::from_pointee(tls_builder.build()));

        let mut last_cert_modified = Self::get_modified_time(&cert_store.certificate_path).await?;
        let mut last_key_modified = Self::get_modified_time(&cert_store.private_key_path).await?;
        let tls_acceptor_arc_clone = Arc::clone(&tls_acceptor_arc);

        let certificate_bundle = Arc::new(ArcSwap::from_pointee(
            CertificateBundle::from_cert_store(&cert_store).await?,
        ));

        let certificate_bundle_clone = Arc::clone(&certificate_bundle);

        tokio::spawn(async move {
            let mut certs_changed_watch = tokio::time::interval(tokio::time::Duration::from_secs(
                CERT_FILES_CHANGE_WATCH_INTERVAL,
            ));

            loop {
                certs_changed_watch.tick().await;
                if let (Ok(cert_m), Ok(key_m)) = (
                    Self::get_modified_time(&cert_store.certificate_path).await,
                    Self::get_modified_time(&cert_store.private_key_path).await,
                ) {
                    if !(cert_m > last_cert_modified || key_m > last_key_modified) {
                        continue;
                    }

                    log::info!("Reloading TLS certificates since they have been modified.");

                    match CertificateBundle::from_cert_store(&cert_store).await {
                        Ok(new_bundle) => {
                            match docdb_openssl::create_tls_acceptor(&new_bundle, acceptor_builder)
                            {
                                Ok(new_tls_acceptor) => {
                                    tls_acceptor_arc_clone
                                        .store(Arc::new(new_tls_acceptor.build()));
                                    last_cert_modified = cert_m;
                                    last_key_modified = key_m;
                                    certificate_bundle_clone.store(Arc::new(new_bundle));

                                    log::info!("TLS certificates reloaded.");
                                }
                                Err(e) => log::error!("Failed to create TLS acceptor: {e:?}."),
                            }
                        }
                        Err(e) => log::error!("Failed to reload TLS certificates: {e:?}."),
                    }
                }
            }
        });

        Ok(Self {
            tls_acceptor: tls_acceptor_arc,
            ciphersuite_mapping,
            certificate_bundle,
        })
    }

    /// Returns the current TLS acceptor.
    pub fn tls_acceptor(&self) -> Arc<SslAcceptor> {
        Arc::clone(&self.tls_acceptor.load())
    }

    /// Maps an SSL ciphersuite to a numeric identifier using the configured cipher mapping function.
    ///
    /// This method converts an SSL cipher reference to a numeric code for telemetry,
    /// logging, or other identification purposes. If no cipher mapping function was
    /// provided during construction, returns 0 as a default value.
    ///
    /// # Arguments
    ///
    /// * `ciphersuite` - An optional reference to an SSL cipher from the current connection.
    ///   This is typically obtained from `SslRef::current_cipher()` on an active SSL connection.
    ///
    /// # Returns
    ///
    /// Returns an `i32` representing the numeric identifier for the cipher:
    /// * If a cipher mapping function is configured and a cipher is provided, returns the mapped value
    /// * If no cipher mapping function is configured, returns 0
    /// * If the cipher is `None`, the mapping function receives `None` and may return a default value
    pub fn ciphersuite_to_i32(&self, ciphersuite: Option<&SslCipherRef>) -> i32 {
        self.ciphersuite_mapping
            .map_or(0, |mapping| mapping(ciphersuite.map(SslCipherRef::name)))
    }

    pub fn is_valid_certificate(&self) -> bool {
        let bundle = self.certificate_bundle.load();
        let certificate = bundle.certificate();

        // private key is RSA
        let is_private_rsa = bundle.private_key.id() == Id::RSA;

        // certificate public key is RSA (and obtainable)
        let pubkey = match bundle.certificate.public_key() {
            Ok(k) => k,
            Err(_) => {
                log::error!("Detected invalid certificate. Failed to obtain public key.");
                return false;
            }
        };

        let is_pub_rsa = pubkey.id() == Id::RSA;

        // private key matches the certificate’s public key
        let matches_cert = pubkey.public_eq(&bundle.private_key);

        let result = is_private_rsa && is_pub_rsa && matches_cert;

        if !result {
            match TlsProvider::sha1_thumbprint(&certificate) {
                Ok(thumbprint) => log::error!("Detected invalid certificate. Thumbprint: {thumbprint:?}"),
                Err(e) => log::error!("Detected invalid certificate. Failed to generate certificate. Thumbprint: {e:?}"),
            }
        }

        result
    }

    fn sha1_thumbprint(cert: &X509) -> Result<String> {
        let der = cert.to_der()?;
        let mut h = Hasher::new(MessageDigest::sha1())?;
        h.update(&der)?;
        let digest = h.finish()?;
        Ok(hex::encode_upper(digest))
    }

    async fn get_modified_time(path: &str) -> Result<std::time::SystemTime> {
        let metadata = tokio::fs::metadata(path).await?;
        Ok(metadata.modified()?)
    }
}
